"""Pass 系统基础类

提供 Pass 系统的基础抽象类和管道实现。
"""

from abc import ABC, abstractmethod
from typing import Any, List, TYPE_CHECKING
from enum import Enum


if TYPE_CHECKING:
    from ..compiler import MCPyCompiler


class PassStage(Enum):
    """Pass 执行阶段"""

    PRE_ANALYSIS = "pre_analysis"  # 分析前预处理
    TYPE_ANALYSIS = "type_analysis"  # 类型分析阶段
    POST_ANALYSIS = "post_analysis"  # 分析后处理
    OPTIMIZATION = "optimization"  # 优化阶段
    CODE_GEN = "code_gen"  # 代码生成阶段


PASS_STAGES = [
    PassStage.PRE_ANALYSIS,
    PassStage.TYPE_ANALYSIS,
    PassStage.POST_ANALYSIS,
    PassStage.OPTIMIZATION,
    PassStage.CODE_GEN,
]


class PassContext:
    """Pass 上下文"""

    def __init__(self, compiler: "MCPyCompiler"):
        self.compiler = compiler


class BasePass(ABC):
    """Pass 基类

    支持五个执行阶段：
    - PRE_ANALYSIS: 分析前预处理
    - TYPE_ANALYSIS: 类型分析阶段
    - POST_ANALYSIS: 分析后处理
    - OPTIMIZATION: 优化阶段
    - CODE_GEN: 代码生成阶段

    子类应该重载 run 方法实现具体逻辑。
    """

    # Pass 名称，子类应该设置此属性
    name: str = ""

    # Pass 执行阶段，子类应该设置此属性
    stage: PassStage = PassStage.POST_ANALYSIS

    @abstractmethod
    def run(self, ir_module: Any, context: PassContext) -> Any:
        """运行 Pass 的核心逻辑

        Args:
            ir_module: IR 模块

        Returns:
            处理后的 IR 模块
        """
        pass

    def should_run_in_stage(self, stage: PassStage) -> bool:
        """判断是否应该在指定阶段运行

        Args:
            stage: 目标执行阶段

        Returns:
            如果应该在该阶段运行则返回 True
        """
        return self.stage == stage


def run_passes_in_stage(
    passes: List[BasePass], stage: PassStage, ir_module: Any, context: PassContext
) -> Any:
    """运行指定阶段的 Pass

    Args:
        passes: Pass 列表
        stage: 执行阶段
        ir_module: 输入的 IR 模块
        context: Pass 上下文

    Returns:
        处理后的 IR 模块
    """
    current_module = ir_module

    # 筛选出应该在该阶段运行的 Pass
    stage_passes = [p for p in passes if p.should_run_in_stage(stage)]

    # 如果没有 Pass 需要运行，直接返回
    if not stage_passes:
        return current_module

    # 按顺序执行 Pass
    for pass_instance in stage_passes:
        current_module = pass_instance.run(current_module, context)

    return current_module


def run_extensions_in_stage(
    stage: PassStage, timing: str, ir_module: Any, extension_registry, compiler=None
) -> Any:
    """运行指定阶段和时机的 Extensions

    Args:
        stage: 执行阶段
        timing: 执行时机 ("pre" 或 "post")
        ir_module: 输入的 IR 模块
        extension_registry: 扩展注册表
        compiler: 编译器实例（可选）

    Returns:
        处理后的 IR 模块
    """

    return extension_registry.run_extensions(stage, timing, ir_module, compiler)
